#!/usr/bin/env python3
"""
multi_mode_manager.py:
管理多地图环境下的导航，包含子进程启动/停止、目标点中继、可视化和Dijkstra路径规划
依赖：ROS Noetic, Python3, rospy, tf, subprocess, threading, yaml, nav_msgs, move_base_msgs, rospkg
"""
import rospy
import yaml
import heapq
import numpy as np
import tf
import tf2_ros
import subprocess
import threading
import rospkg
from geometry_msgs.msg import PoseWithCovarianceStamped, Point, Quaternion, PoseStamped, TransformStamped
from std_msgs.msg import Header
from geometry_msgs.msg import PointStamped
from nav_msgs.msg import OccupancyGrid
from tf.transformations import euler_from_quaternion

class DijkstraPlanner:
    def __init__(self, graph):
        self.nodes = {n['block_name']: n for n in graph['nodes']}
        self.adj = {}
        for e in graph['edges']:
            frm, to = e['from'], e['to']
            self.adj.setdefault(frm, []).append((to, e))

    def shortest_path(self, start, goal):
        dist = {n: float('inf') for n in self.nodes}
        prev = {}
        dist[start] = 0
        pq = [(0, start)]
        while pq:
            d, u = heapq.heappop(pq)
            if u == goal:
                break
            if d > dist[u]: continue
            for v, edge in self.adj.get(u, []):
                alt = d + 1
                if alt < dist[v]:
                    dist[v] = alt
                    prev[v] = (u, edge)
                    heapq.heappush(pq, (alt, v))
        if goal not in prev and start != goal:
            return []
        path = []
        node = goal
        while node != start:
            u, edge = prev[node]
            path.append((u, node, edge))
            node = u
        path.reverse()
        return path


class MapVisualizer:
    def __init__(self, graph, pkg_path):
        self.graph = graph
        self.pkg_path = pkg_path
        self.transforms = []
        self.map_processes = []
        self.active_node = None
        # 加载并计算各子地图相对 merged_map 的静态变换
        self.load_and_arrange_maps()

        # 发布静态变换： merged_map -> 子地图_frame
        static_broadcaster = tf2_ros.StaticTransformBroadcaster()
        static_ts = []
        now = rospy.Time.now()
        for ns, (tx, ty, tz), (rx, ry, rz, rw), _ in self.transforms:
            t = TransformStamped()
            t.header.stamp = now
            t.header.frame_id = 'merged_map'
            t.child_frame_id = ns  # 原为 f"{ns}/map"
            t.transform.translation.x = tx
            t.transform.translation.y = ty
            t.transform.translation.z = tz
            t.transform.rotation.x = rx
            t.transform.rotation.y = ry
            t.transform.rotation.z = rz
            t.transform.rotation.w = rw
            static_ts.append(t)
        static_broadcaster.sendTransform(static_ts)

        # 动态发布 map -> active 子地图_frame
        self.tf_broadcaster = tf.TransformBroadcaster()
        threading.Thread(target=self.publish_active_transform, daemon=True).start()

    def load_and_arrange_maps(self):
        margin = 20.0
        cumulative_y = 0.0
        for node in self.graph['nodes']:
            block = node['block_name']
            raw = node['map_file']
            if raw.startswith('$(find'):
                map_file = raw.replace('$(find navigation_pkg)', self.pkg_path)
            else:
                map_file = raw

            # 启动 map_server，frame_id 设为子地图_frame（去除 "/map"）
            map_ns = f"{block}_map"
            proc = subprocess.Popen([
                'rosrun', 'map_server', 'map_server', map_file,
                '__ns:=' + map_ns,
                '_frame_id:=' + map_ns
            ])
            self.map_processes.append(proc)
            topic = f"/{map_ns}/map"
            rospy.loginfo(f"Waiting for map on {topic}...")
            msg = rospy.wait_for_message(topic, OccupancyGrid)
            info = msg.info

            # 计算地图尺寸与边界
            width_m = info.width * info.resolution
            height_m = info.height * info.resolution
            tx = 0.0
            ty = cumulative_y
            tz = 0.0
            x1, y1 = tx + info.origin.position.x, ty + info.origin.position.y
            x2, y2 = x1 + width_m , y1 + height_m 
            self.transforms.append((map_ns, (tx, ty, tz), (0.0, 0.0, 0.0, 1.0), [x1, x2, y1, y2]))
            cumulative_y += height_m + margin
        rospy.loginfo(str(self.transforms))

    def set_active_node(self, node_name):
        self.active_node = node_name

    def publish_active_transform(self):
        rate = rospy.Rate(10)
        while not rospy.is_shutdown():
            now = rospy.Time.now()
            for ns, _, rot, _ in self.transforms:
                if ns == f"{self.active_node}_map":
                    # 动态发布 map -> 子地图_frame
                    self.tf_broadcaster.sendTransform(
                        (0.0, 0.0, 0.0),
                        rot,
                        now,
                        'map',
                        ns  
                    )
            rate.sleep()


class NavigationManager:
    def __init__(self):
        rospy.init_node('multi_mode_manager')
        path = rospy.get_param('~node_graph_path')
        with open(path, 'r') as f:
            graph = yaml.safe_load(f)
        self.graph = graph
        rospack = rospkg.RosPack()
        self.pkg_path = rospack.get_path('navigation_pkg')
        self.planner = DijkstraPlanner(graph)
        self.visualizer = MapVisualizer(graph, self.pkg_path)
        self.current_node = None
        self.current_proc = None
        self.tf_listener = tf.TransformListener()
        self.thresh_position = 0.5
        self.thresh_pose = 0.3

        # 异步导航计划: segments 列表 和 当前索引
        self.nav_plan = []  # 每一项 [last_goal,node_name, init_pose, next_goal] , 最后一项 "Done"
        self.nav_goal = None
        self.current_goal = [] # x,y,z,roll,pitch,yaw
        self.timer = rospy.Timer(rospy.Duration(0.1), self.timer_callback)

        rospy.Subscriber('/clicked_point', PointStamped, self.click_callback)
        rospy.Subscriber('/nav_goal', PoseStamped, self.nav_goal_callback)
        self.goal_pub = rospy.Publisher('/move_base_simple/goal', PoseStamped, queue_size=1)

        # 初始化第一个节点
        first = graph['nodes'][0]
        self.switch_to_node(first['block_name'], first['default_pose'])

    def timer_callback(self, event):
        # 如果无导航目标点则直接腿粗
        if self.nav_goal == None:
            return 
        # 如果有导航目标点但是没有导航计划则生成导航计划
        if len(self.nav_plan)<=0 :
            self.make_nav_plan()
            
        # 如果只剩一项，说明已经到达目标层，则直接发布目标
        if len(self.nav_plan)==1:
            self.publish_goal(self.nav_goal)
            self.nav_goal = None
            return
        # 有导航目标点和导航计划则检测自己的位置,如果没有到达就继续运行
        if not self.check_arrival():
            return 
        # 已经到达目标则进行地图切换并重新发布新的目标坐标
        last_goal, node_name,init_pose,goal_pose=self.nav_plan.pop(0)
        self.switch_to_node(node_name,init_pose)
        self.publish_goal(goal_pose)
        self.current_goal = goal_pose       

    def check_arrival(self):
        try:
            # 获取从 map 到 body 的 tf 变换
            self.tf_listener.waitForTransform("map", "body", rospy.Time(0), rospy.Duration(1.0))
            transform = self.tf_listener.lookupTransform("map", "body", rospy.Time(0))

            # 提取位置信息 (x, y, z)
            x = transform[0][0]  # 索引修正为从 lookupTransform 返回的 tuple 中提取数据
            y = transform[0][1]
            z = transform[0][2]  # 修正了原代码中重复使用 y 的问题

            # 提取方向信息 (yaw)，通过四元数转换为欧拉角
            orientation = transform[1]  # 索引修正为从 lookupTransform 返回的 tuple 中提取数据
            quaternion = [orientation[0], orientation[1], orientation[2], orientation[3]]
            roll, pitch, yaw = euler_from_quaternion(quaternion)  # 假设 euler_from_quaternion 是外部定义的函数

            # 当前位姿
            now = [x, y, z, roll % (2*3.14159), pitch % (2*3.14159), yaw % (2*3.14159)]
            now_position = [x,y,z]
            now_pose = [roll,pitch,yaw]
            target_position = self.current_goal[:3]
            target_pose = self.current_goal[3:]

            # diff = np.linalg.norm(np.array(now) - np.array(self.current_goal))
            diff_position = np.linalg.norm(np.array(now_position) - np.array(target_position))
            
            # 求位姿差
            diff_pose = np.abs(np.array(now_pose) - np.array(target_pose))
            for i in range(len(diff_pose)):
                diff_pose[i] = diff_pose[i] if diff_pose[i] < np.pi else 2*3.1415926-diff_pose[i]
            diff_pose=np.linalg.norm(diff_pose)
            
            rospy.loginfo(f"Position difference: {diff_position}, Pose difference: {diff_pose} , current pose is {now}")
            # 判断是否到达目标点
            if diff_position < self.thresh_position and diff_pose < self.thresh_pose:
                return True
            return False

        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
            rospy.logwarn(f"TF lookup failed: {e}")
            return False
    
    def make_nav_plan(self):
        from_node = self.current_node
        to_node = self.nav_goal[-1]
        # 使用 Dijkstra 算法完成路径规划
        path = self.planner.shortest_path(from_node,to_node)
        # 路径转化为需要的格式
        # edge = {
        # "from" : node_name,
        # "to" : node2_name,
        # "switch_point_at_from" : [x y z roll pitch yaw]
        # "switch_point_at_to" :  [x y z roll pitch yaw]
        # }
        self.nav_plan =[[edge['switch_point_at_from'],edge["to"],edge['switch_point_at_to']] for f,t,edge in path]
        for i in range(len(self.nav_plan)-1):
            self.nav_plan[i].append(self.nav_plan[i+1][0])
        if len(self.nav_plan)!=0:
            self.nav_plan[-1].append(self.nav_goal)
        self.nav_goal.pop()
        # 补充最后一位
        self.nav_plan.append("Done")
        # 发布一次
        if len(self.nav_plan)>1:
            self.current_goal = self.nav_plan[0][0]
            self.publish_goal(self.current_goal)
        rospy.loginfo(f"Plan have been made {str(self.nav_plan)}")
    
    def publish_goal(self,goal):
        rospy.loginfo(goal)
        if len(goal)==7:
            goal.pop()
        x,y,z,roll,pitch,yaw = goal
        msg = PoseStamped()
        msg.header = Header(stamp=rospy.Time.now(), frame_id="map")
        msg.pose.position = Point(x,y,z)
        quat = tf.transformations.quaternion_from_euler(roll, pitch, yaw)
        msg.pose.orientation = Quaternion(*quat)
        # 发布
        self.goal_pub.publish(msg)
        rospy.loginfo(f"Goal Published : {goal}")
    
    def nav_goal_callback(self, msg: PoseStamped):
        try:
            # 提取目标帧的命名空间
            ns = msg.header.frame_id.split('/')[0]
            target = ns.replace('_map', '')  # 去掉 '_map' 后缀

            # 修改 frame_id 为目标地图
            msg.header.frame_id = target + '_map'

            # 将目标位姿从目标地图坐标系转换到全局 'map' 坐标系
            mapped = self.tf_listener.transformPose('map', msg)

            # 检测目标是否存在于图中的节点
            flag = False
            for node in self.graph["nodes"]:
                if node["block_name"] == target:
                    flag = True
                    break

            if not flag:
                # 如果目标不在图中，清空导航目标
                self.nav_goal = None
                rospy.logwarn(f"Target '{target}' not found in the graph.")
            else:
                # 提取位置信息 (x, y, z)
                x = mapped.pose.position.x
                y = mapped.pose.position.y
                z = mapped.pose.position.z

                # 提取方向信息 (roll, pitch, yaw)，通过四元数转换为欧拉角
                orientation = mapped.pose.orientation
                quaternion = [orientation.x, orientation.y, orientation.z, orientation.w]
                roll, pitch, yaw = tf.transformations.euler_from_quaternion(quaternion)

                # 设置导航目标
                self.nav_goal = [x, y, z, roll, pitch, yaw, target]
                self.nav_plan = []

        except Exception as e:
            rospy.logerr(f'Nav goal callback error: {e}')

    def click_callback(self, msg):
        try:
            block, pose = self.transform_to_block(msg)
            if block is None:
                rospy.logerr('Clicked point is outside all known maps.')
                self.nav_goal = None
                return
            self.nav_goal = pose
            self.nav_goal.append(block)
            self.nav_plan = []
        except Exception as e:
            rospy.logerr(f'Click callback error: {e}')

    def transform_to_block(self, point_msg):
        rospy.loginfo("Recv goal, processing")
        
        # 确保时间戳为最新
        point_msg.header.stamp = rospy.Time(0)
        
        # 等待变换变为可用，增加超时时间
        try:
            self.tf_listener.waitForTransform("merged_map", point_msg.header.frame_id, rospy.Time(0), rospy.Duration(1.0))
        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException) as e:
            rospy.logerr(f"Wait for transform failed: {e}")
            return None, None
        
        rospy.loginfo(f"Transforming goal from {point_msg.header.frame_id} to merged_map")
        
        # 执行坐标转换
        try:
            merged_pt = self.tf_listener.transformPoint("merged_map", point_msg)
        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException) as e:
            rospy.logerr(f"Transform point failed: {e}")
            return None, None
        x, y = merged_pt.point.x, merged_pt.point.y
        
        rospy.loginfo("Trying to find out which map goal belongs to")
        # 使用边界坐标判断点击点所属地图
        for ns, trans, rot, bounds in self.visualizer.transforms:
            x1, x2, y1, y2 = bounds
            if x1 <= x and x <= x2 and y1 <= y and y <= y2:
                # 计算相对于地图原点的坐标
                rel_x = x - trans[0]
                rel_y = y - trans[1]
                block_name = ns.replace("_map", "")       
                rospy.loginfo(f"Navigation Goal belong to {block_name}")
                return block_name, [rel_x, rel_y, 0  ,0,0,0]
        return None, None

    def switch_to_node(self, node_name, pose6):
        # 完全终止旧进程
        if self.current_proc:
            # 使用更安全的终止方式
            self.current_proc.terminate()
            self.current_proc.wait()  # 等待进程结束
            # 精确终止导航节点（原pkill命令过于危险）
            subprocess.call(['rosnode', 'kill', self.current_node])  # 根据实际节点名调整
            rospy.loginfo(f"Stopped navigation on {self.current_node}")
            rospy.sleep(3.0)  # 确保节点完全终止            
            
        # 启动新进程（保持原有逻辑）
        raw = self.get_node(node_name)['launch_file']
        if raw.startswith('$(find'):
            launch_file = raw.replace('$(find navigation_pkg)', self.pkg_path)
        else:
            launch_file = raw
        
        self.current_proc = subprocess.Popen(['roslaunch', launch_file])
        self.current_node = node_name
        self.visualizer.set_active_node(node_name)
        rospy.loginfo(f"Launched navigation on {node_name}")
        pos = pose6[:3]
        yaw = pose6[5]
        self.ready_to_navigate=False
        self.publish_initial_pose(pos, (0, 0, yaw))

    def get_node(self, name):
        for n in self.graph['nodes']:
            if n['block_name'] == name:
                return n
        return None

    def publish_initial_pose(self, position, euler):
        pub = rospy.Publisher('/initialpose', PoseWithCovarianceStamped, queue_size=2)
        rospy.sleep(15.0)
        quat = tf.transformations.quaternion_from_euler(*euler)
        init = PoseWithCovarianceStamped()
        init.header = Header(stamp=rospy.Time.now(), frame_id='map')
        init.pose.pose.position = Point(*position)
        init.pose.pose.orientation = Quaternion(*quat)
        pub.publish(init)
        rospy.loginfo('---------------------------------------------------------------')        
        rospy.loginfo('---------------------------------------------------------------')
        rospy.loginfo('Initial pose published')
        rospy.loginfo('---------------------------------------------------------------')
        rospy.loginfo('---------------------------------------------------------------')
        self.ready_to_navigate = True


if __name__ == '__main__':
    try:
        manager = NavigationManager()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
